package io.github.quickmsg.interate.sharedstrategy;

import cn.hutool.core.collection.ListUtil;
import org.apache.commons.lang3.tuple.ImmutableTriple;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

public class ConsistentHashStrategy implements SharedStrategy{
    @Override
    public ImmutableTriple<String, String, String> searchSharedSubscribe(Set<ImmutableTriple<String, String, String>> shareSubscribes, String consistentId) {
        return new ConsistentHashSelector<>(ListUtil.toList(shareSubscribes), 10).select(consistentId);
    }

    private static final class ConsistentHashSelector<T> {

        private final SortedMap<Long, T> virtualShareSubscribes = new TreeMap<>();
        private final HashFunction hashFunction = new MD5Hash();

        ConsistentHashSelector(List<T> shareSubscribes, int virtualNodes) {
            for (T shareSubscribe : shareSubscribes) {
                for (int i = 0; i < virtualNodes; i++) {
                    virtualShareSubscribes.put(hashFunction.hash(shareSubscribe.toString() + i), shareSubscribe);
                }
            }
        }

        public T select(String objectKey) {
            SortedMap<Long, T> tailMap = virtualShareSubscribes.tailMap(hashFunction.hash(objectKey));
            Long nodeHashVal = tailMap.isEmpty() ? virtualShareSubscribes.firstKey() : tailMap.firstKey();
            return virtualShareSubscribes.get(nodeHashVal);
        }
    }

    private static class MD5Hash implements HashFunction {
        MessageDigest instance;
        public MD5Hash() {
            try {
                instance = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                throw new IllegalStateException(e.getMessage(), e);
            }
        }

        @Override
        public long hash(String key) {
            instance.reset();
            instance.update(key.getBytes());
            byte[] digest = instance.digest();
            long h = 0;
            for (int i = 0; i < 4; i++) {
                h <<= 8;
                h |= ((int) digest[i]) & 0xFF;
            }
            return h;
        }
    }

    /**
     * Hash String to long value
     */
    public interface HashFunction {
        long hash(String key);
    }
}
